# LeNet分为卷积层块和全连接层块两个部分。
# 卷积层块里的基本单位是卷积层后接最大池化层：
# 在卷积层块中，每个卷积层都使用5×5的窗口，使用sigmoid激活函数。
# 第一个卷积层输出通道数为6，第二个卷积层输出通道数则增加到16。
# 卷积层块的两个最大池化层的窗口形状均为2×2，且步幅为2。
# 池化窗口在输入上每次滑动所覆盖的区域互不重叠。
# 卷积层块的输出形状为(批量大小, 通道, 高, 宽)。
# 当卷积层块的输出传入全连接层块时,全连接层块含3个全连接层。
# 它们的输出个数分别是120、84和10，其中10为输出的类别个数。


# 通过Sequential类来实现LeNet模型。

import time
import torch
from torch import nn, optim

from basic_knowledge import d2lzh_pytorch as d2l
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 6, 5),  # in_channels, out_channels, kernel_size
                                 # 一张图片，六个核，窗口5×5
                                 # 28×28变24×24
            nn.Sigmoid(),
            nn.MaxPool2d(2, 2),  # kernel_size, stride
                                # 输出6通道，特征图大小减半，24×24变12×12
            nn.Conv2d(6, 16, 5),   # 输入6通道，输出16通道，窗口5×5，12×12变8×8
            nn.Sigmoid(),
            nn.MaxPool2d(2, 2)  # 输出16通道，特征图大小减半，8×8变4×4
        )
        self.fc = nn.Sequential(
            nn.Linear(16*4*4, 120),  # flaten输入向量，输出120
            nn.Sigmoid(),
            nn.Linear(120, 84),  # 输出84
            nn.Sigmoid(),
            nn.Linear(84, 10)  # 输出10
        )

    def forward(self, img):
        feature = self.conv(img)
        output = self.fc(feature.view(img.shape[0], -1))
        return output


net = LeNet()
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size=batch_size)


# 学习率采用0.001，训练算法使用Adam算法，损失函数使用交叉熵损失函数。
lr, num_epochs = 0.001, 5
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
d2l.train_ch5(net, train_iter, test_iter, batch_size, optimizer, device, num_epochs)